package cn.zhage.aimin.ai.config;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/*
 *
 * @author 渣哥
 */
@Configuration
@RequiredArgsConstructor
public class ChatClientConfig {

    private final ChatStorageMemory chatMemory;

    private final VectorStore vectorStore;

    @Bean
    public ChatClient openAiChatClient(OllamaChatModel chatModel, ToolCallbackProvider provider) {
        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(chatMemory).build();

        var qaAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(SearchRequest.builder().similarityThreshold(0.4d).topK(2).build())
                .build();

        return ChatClient.builder(chatModel)
                .defaultToolCallbacks(provider)
                .defaultAdvisors(
                        new SimpleLoggerAdvisor(),
                        chatMemoryAdvisor,
                        qaAdvisor
                )
                .build();
    }

}
